from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--model_type', type=str, default='us_v10')
parser.add_argument('--execute_schemes', type=str, help='dataset:scheme id;...')
parser.add_argument('--seed', type=int, default=42)
parser.add_argument('--execute_threads', type=int, default=4)
parser.add_argument('--few_shot', action='store_true')
parser.add_argument('--one_shot', action='store_true')
parser.add_argument('--zero_shot', action='store_true')
parser.add_argument('--insert_start_token', action='store_true')
parser.add_argument('--only_calculate_last', action='store_true')
parser.add_argument('--test_dev', action='store_true')

args = parser.parse_args()

model_type = args.model_type
few_shot = args.few_shot
one_shot = args.one_shot
zero_shot = args.zero_shot
insert_start_token = args.insert_start_token
only_calculate_last = args.only_calculate_last
test_dev = args.test_dev
execute_threads = int(args.execute_threads)

assert model_type in {
    'pangu-2.6b', 'pangu-13b-fp16', 'us-v9', 'us-v10', 'us-v11', 'us-002-v3'
}


def get_schemes() -> dict:
    from backend.experiment.clue.afqmc.scheme0.implement import \
        AFQMCImplement as AFQMCImplementS0
    from backend.experiment.clue.afqmc.scheme1.implement import \
        AFQMCImplement as AFQMCImplementS1
    from backend.experiment.clue.afqmc.scheme2.implement import \
        AFQMCImplement as AFQMCImplementS2
    from backend.experiment.clue.afqmc.scheme3.implement import \
        AFQMCImplement as AFQMCImplementS3
    from backend.experiment.clue.afqmc.scheme4.implement import \
        AFQMCImplement as AFQMCImplementS4
    from backend.experiment.clue.afqmc.scheme5.implement import \
        AFQMCImplement as AFQMCImplementS5
    from backend.experiment.clue.afqmc.scheme6.implement import \
        AFQMCImplement as AFQMCImplementS6
    from backend.experiment.clue.afqmc.scheme7.implement import \
        AFQMCImplement as AFQMCImplementS7
    from backend.experiment.clue.c3.scheme0.implement import \
        C3Implement as C3ImplS0
    from backend.experiment.clue.c3.scheme1.implement import \
        C3Implement as C3ImplS1
    from backend.experiment.clue.chid.scheme0.implement import \
        CHIDImplement as CHIDImplS0
    from backend.experiment.clue.chid.scheme1.implement import \
        CHIDImplement as CHIDImplS1
    from backend.experiment.clue.chid.scheme2.implement import \
        CHIDImplement as CHIDImplS2
    from backend.experiment.clue.chid.scheme3.implement import \
        CHIDImplement as CHIDImplS3
    from backend.experiment.clue.cluewsc2020.scheme0.implement import \
        CLUEWSC2020Implement as CLUEWSC2020ImplS0
    from backend.experiment.clue.cmnli.scheme0.implement import \
        CMNLI2020Implement as CMNLI2020ImplS0
    from backend.experiment.clue.cmrc2018.scheme0.implement import \
        CMRC2018Implement as CMRC2018ImplS0
    from backend.experiment.clue.csl.scheme0.implement import \
        CSLImplement as CSLImplS0
    from backend.experiment.clue.csl.scheme1.implement import \
        CSLImplement as CSLImplS1
    from backend.experiment.clue.drcd.scheme0.implement import \
        DRCDImplement as DRCDImplS0
    from backend.experiment.clue.iflytek.scheme0.implement import \
        IFlyTekImplement as IFlyTekImplS0
    from backend.experiment.clue.ocnli.scheme0.implement import \
        OCNLIImplement as OCNLIImplS0
    from backend.experiment.clue.ocnli.scheme1.implement import \
        OCNLIImplement as OCNLIImplS1
    from backend.experiment.clue.tnews.scheme0.implement import \
        TNewsImplement as TNewsImplS0
    from backend.experiment.raw_text_perplexity.implement import \
        RawTextPerplexityImplement as RawTextPerplexityImplementS0
    from backend.experiment.raw_text_perplexity.clue_oscar.\
        clue_oscar_implement import \
        CLUEOscarTextPerplexityImplement as RawTextPerplexityImplementS1

    return {
        'clue.afqmc:0': AFQMCImplementS0,
        'clue.afqmc:1': AFQMCImplementS1,
        'clue.afqmc:2': AFQMCImplementS2,
        'clue.afqmc:3': AFQMCImplementS3,
        'clue.afqmc:4': AFQMCImplementS4,
        'clue.afqmc:5': AFQMCImplementS5,
        'clue.afqmc:6': AFQMCImplementS6,
        'clue.afqmc:7': AFQMCImplementS7,
        'clue.c3:0': C3ImplS0,
        'clue.c3:1': C3ImplS1,
        'clue.chid:0': CHIDImplS0,
        'clue.chid:1': CHIDImplS1,
        'clue.chid:2': CHIDImplS2,
        'clue.chid:3': CHIDImplS3,
        'clue.cluewsc2020:0': CLUEWSC2020ImplS0,
        'clue.cmnli:0': CMNLI2020ImplS0,
        'clue.cmrc2018:0': CMRC2018ImplS0,
        'clue.csl:0': CSLImplS0,
        'clue.csl:1': CSLImplS1,
        'clue.drcd:0': DRCDImplS0,
        'clue.iflytek:0': IFlyTekImplS0,
        'clue.ocnli:0': OCNLIImplS0,
        'clue.ocnli:1': OCNLIImplS1,
        'clue.tnews:1': TNewsImplS0,
        'rawtextppl:0': RawTextPerplexityImplementS0,
        'rawtextppl:1': RawTextPerplexityImplementS1
    }


def execute_schemes():
    schemes = get_schemes()
    return [schemes[each] for each in args.execute_schemes.split(',')]


def shots_prompt_modes() -> list:
    from backend.experiment.framework.shots_prompt.implement import ShotMode
    res = []
    if args.few_shot:
        res.append((ShotMode.FEW_SHOT, 'few_shot'))
    if args.one_shot:
        res.append((ShotMode.ONE_SHOT, 'one_shot'))
    if args.zero_shot:
        res.append((ShotMode.ZERO_SHOT, 'zero_shot'))
    return res


seed = int(args.seed)
